# -*- coding: utf-8 -*-
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QFileDialog, QMessageBox
import os
import numpy as np
from ultralytics import YOLO


class ImageViewer(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.setup_ui()
        self.original_image = None # 存储原始图像
        self.image_path = None # 图像文件路径

        # 加载YOLO模型
        self.model = YOLO('best.pt')

    def setup_ui(self):
        self.setWindowTitle("图片查看器")
        self.resize(800, 600)

        # 中央部件
        central_widget = QtWidgets.QWidget()
        self.setCentralWidget(central_widget)

        # 图像显示区域
        self.image_label = QtWidgets.QLabel()
        self.image_label.setAlignment(QtCore.Qt.AlignCenter)
        self.image_label.setStyleSheet("""
            QLabel {
                background-color: #f0f0f0;
                border: 2px dashed #aaa;
                qproperty-alignment: 'AlignCenter';
            }
        """)
        self.image_label.setText("图片显示区域")

        # 按钮区域
        self.upload_btn = QtWidgets.QPushButton("上传图片")
        self.recognize_btn = QtWidgets.QPushButton("识别")

        # 布局
        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.image_label, 1)

        btn_layout = QtWidgets.QHBoxLayout()
        btn_layout.addWidget(self.upload_btn)
        btn_layout.addWidget(self.recognize_btn)
        layout.addLayout(btn_layout)

        central_widget.setLayout(layout)

        # 连接信号
        self.upload_btn.clicked.connect(self.upload_image)
        self.recognize_btn.clicked.connect(self.run_detection)

    def upload_image(self):
        """上传并显示图片"""
        file_name, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        ) # 弹出文件选择对话框

        if file_name:
            self.image_path = file_name
            self.original_image = QtGui.QImage(file_name) # 加载为QImage
            if self.original_image.isNull():
                QMessageBox.warning(self, "错误", "无法加载图片！")
                return
            self.display_image(self.original_image)# 显示图像

    def run_detection(self):
        """使用YOLO进行识别并显示结果"""
        if not hasattr(self, 'image_path') or not self.image_path:
            QMessageBox.warning(self, "错误", "请先上传图片！")
            return

        try:
            # 使用YOLO进行预测
            results = self.model.predict(self.image_path)

            # 获取带标注的结果图像
            result_img = results[0].plot()  # 获取带检测框的图像(numpy数组)

            # 将numpy数组转换为QImage
            height, width, channel = result_img.shape
            bytes_per_line = 3 * width
            qimage = QtGui.QImage(result_img.data, width, height, bytes_per_line, QtGui.QImage.Format_RGB888)

            # 显示结果图像
            self.display_image(qimage)

        except Exception as e:
            QMessageBox.critical(self, "识别错误", f"发生错误: {str(e)}")

    def display_image(self, image):
        """显示图片（自动缩放）"""
        if isinstance(image, QtGui.QImage):
            scaled = image.scaled(
                self.image_label.width() - 20,
                self.image_label.height() - 20,
                QtCore.Qt.KeepAspectRatio,
                QtCore.Qt.SmoothTransformation
            )
            self.image_label.setPixmap(QtGui.QPixmap.fromImage(scaled))


if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)


    window = ImageViewer()
    window.show()
    sys.exit(app.exec_())